# -*- coding: utf-8 -*-

my_data = [['slashdot', 'USA', 'yes', 18, 'None'],
        ['google', 'France', 'yes', 23, 'Premium'],
        ['digg', 'USA', 'yes', 24, 'Basic'],
        ['kiwitobes', 'France', 'yes', 23, 'Basic'],
        ['google', 'UK', 'no', 21, 'Premium'],
        ['(direct)', 'New Zealand', 'no', 12, 'None'],
        ['(direct)', 'UK', 'no', 21, 'Basic'],
        ['google', 'USA', 'no', 24, 'Premium'],
        ['slashdot', 'France', 'yes', 19, 'None'],
        ['digg', 'USA', 'no', 18, 'None'],
        ['google', 'UK', 'no', 18, 'None'],
        ['kiwitobes', 'UK', 'no', 19, 'None'],
        ['digg', 'New Zealand', 'yes', 12, 'Basic'],
        ['slashdot', 'UK', 'no', 21, 'None'],
        ['google', 'UK', 'yes', 18, 'Basic'],
        ['kiwitobes', 'France', 'yes', 19, 'Basic']]

def divideset(rows, column, value):
    split_function = None
    if isinstance(value, int) or isinstance(value, float):
        split_function = lambda row:row[column] >= value
    else:
        split_function = lambda row:row[column] == value

    set1 = [row for row in rows if split_function(row)]
    set2 = [row for row in rows if not split_function(row)]
    return (set1, set2)


def uniquecounts(rows):
    results = {}
    for row in rows:
        r = row[len(row) - 1]
        if r not in results:results[r] = 0
        results[r] += 1
    return results


def giniimpurity2(rows):
    total = len(rows)
    counts = uniquecounts(rows)
    p = {}    
    for k3 in counts:p[k3] = float(counts[k3]) / total 
    return 1 - sum(pp * pp for pp in p.values())
    

# 基尼不纯度
def giniimpurity(rows):
    total = len(rows)
    counts = uniquecounts(rows)
    imp = 0
    for k1 in counts:
        p1 = float(counts[k1]) / total
        for k2 in counts:
            if k1 == k2:continue
            p2 = float(counts[k2]) / total
            imp += p1 * p2
    return imp

# 熵
def entropy(rows):
    from math import log
    log2 = lambda x:log(x) / log(2)
    results = uniquecounts(rows)
    ent = 0.0
    for r in results.keys():
        p = float(results[r]) / len(rows)
        ent = ent - p * log2(p)
    return ent

def buildtree(rows, scoref=entropy):
    if len(rows) == 0: return decisionnode()
    current_score = scoref(rows)
    
    best_gain = 0.0
    best_criteria = None
    best_set = None
    
    column_count = len(rows[0]) - 1
    
    for col in range(0, column_count):
        column_values = {}
        for row in rows:
            column_values[row[col]] = 1
        for value in column_values.keys():
            (set1, set2) = divideset(rows, col, value)
            p = float(len(set1)) / len(rows)
            # 信息增益,熵代表集合的混亂程度，两个集合熵的加权越大，代表越混乱
            # 此处gain越大，表示这两个结合越有序
            gain = current_score - p * scoref(set1) - (1 - p) * scoref(set2)
            if gain > best_gain and len(set1) > 0 and len(set2) > 0:
                best_gain = gain
                best_criteria = (col, value)
                best_sets = (set1, set2)
      
                   
    if best_gain > 0:
        trueBranch = buildtree(best_sets[0])
        falseBranch = buildtree(best_sets[1])
        return decisionnode(col=best_criteria[0], value=best_criteria[1],
                            tb=trueBranch, fb=falseBranch)
    else:
        #results总是一个key
        return decisionnode(results=uniquecounts(rows))
    
def printtree(tree, indent=''):
    if tree.results != None:
        print str(tree.results)
    else:
        print str(tree.col) + ":" + str(tree.value) + '?'
        
        print indent + 'T->',
        printtree(tree.tb, indent + '    ')
        print indent + 'F->',
        printtree(tree.fb, indent + '    ')

def getwidth(tree):
    if tree.tb == None and tree.fb == None:return 1
    return getwidth(tree.tb) + getwidth(tree.fb)

def getdepth(tree):
    if tree.tb == None and tree.fb == None:return 0
    return max(getdepth(tree.tb), getdepth(tree.fb)) + 1

def drawnode(draw, tree, x, y):
    if tree.results == None:
        # Get the width of each branch
        w1 = getwidth(tree.fb) * 100
        w2 = getwidth(tree.tb) * 100
        
        # Determine the total space required by this node
        left = x - (w1 + w2) / 2
        right = x + (w1 + w2) / 2
        
        # Draw the condition string
        draw.text((x - 20, y - 10), str(tree.col) + ':' + str(tree.value), (0, 0, 0))
        
        # Draw links to the branches
        draw.line((x, y, left + w1 / 2, y + 100), fill=(255, 0, 0))
        draw.line((x, y, right - w2 / 2, y + 100), fill=(255, 0, 0))
        
        # Draw the branch nodes
        drawnode(draw, tree.fb, left + w1 / 2, y + 100)
        drawnode(draw, tree.tb, right - w2 / 2, y + 100)
    else:
        txt = ' \n'.join(['%s:%d' % v for v in tree.results.items()])
        draw.text((x - 20, y), txt, (0, 0, 0))

from PIL import Image, ImageDraw

def drawtree(tree, jpeg='tree.jpg'):
    w = getwidth(tree) * 100
    h = getdepth(tree) * 100 + 120
    
    img = Image.new('RGB', (w, h), (255, 255, 255))
    draw = ImageDraw.Draw(img)
    
    drawnode(draw, tree, w / 2, 20)
    img.save(jpeg, 'JPEG')
    
    

class decisionnode:
    def __init__(self, col= -1, value=None, results=None, tb=None, fb=None):
        self.col = col;
        self.value = value;
        self.results = results
        self.tb = tb
        self.fb = fb
        
    

#递归遍历整个树
def classify(observation,tree):
    if tree.results!=None:
        return tree.results
    else:
        v=observation[tree.col]
        branch=None
        if isinstance(v,int) or isinstance(v,float):
            if v>=tree.value:branch=tree.tb
            else:branch=tree.fb
        else:
            if v==tree.value:branch=tree.tb
            else:branch=tree.fb
        return classify(observation,branch)

def prune(tree,mingain):
    if tree.tb.results == None:prune(tree.tb,mingain)
    if tree.fb.results == None:prune(tree.fb,mingain)
    
    if tree.tb.results!=None and tree.fb.results!=None:
        tb,fb=[],[]
        #参考[[3]]*5 = [[3], [3], [3], [3], [3]]
        for v,c in tree.tb.results.items():
            tb+=[[v]]*c
        for v,c in tree.fb.results.items():
            fb+=[[v]]*c
            
        #求出如果合并之后的熵        
        delta=entropy(tb+fb)-(entropy(tb)+entropy(fb))/2
        if delta<mingain:
            tree.tb,tree.fb=None,None
            tree.results=uniquecounts(tb+fb)
            
def mdclassify(observation,tree):
    if tree.results!=None:
        return tree.results;
    
    v=observation[tree.col]
    if v==None:
        #递归返回每个子节点的加权值，将这些值相加得到缺失节点的值
        tr,fr=mdclassify(observation,tree.tb),mdclassify(observation,tree.fb)
        tcount=sum(tr.values())
        fcount=sum(fr.values())
        tw=float(tcount)/(tcount+fcount)
        fw=float(fcount)/(tcount+fcount)
        result={}
        for k,v in tr.items(): result[k]=v*tw
        for k,v in fr.items():
            if k not in result:result[k]=0
            result[k] += v*fw      
        return result
    else:
        if isinstance(v,int) or isinstance(v,float):
            if v>=tree.value: branch=tree.tb
            else: branch=tree.fb
        else:
            if v==tree.value: branch=tree.tb
            else: branch=tree.fb
        return mdclassify(observation,branch)

        
        
        

if __name__ == '__main__':
#    print uniquecounts(my_data)
#    
#    print giniimpurity(my_data)
#    
#    print giniimpurity2(my_data)
#    
#    print entropy(my_data)    
#    
#    set1,set2=dividset(my_data,2,'yes');
#    
#    print set1
#    
#    print set2
#    
#    print entropy(set1)
#    
#    print entropy(set2)
    
    tree = buildtree(my_data);
    
    printtree(tree)
    
#    drawtree(tree)

#    prune(tree,1.0)

#    printtree(tree)

#    print classify(['(direct)', 'USA', 'yes', 5],tree)

    print mdclassify(['google',None,'yes',None],tree)
    
    
    